import os
import cv2

# 文件路径定义
labels_dir = 'datasets/rcs_trainval_0422/labels/val'
images_dir = 'datasets/rcs_trainval_0422/images/val'
output_dir = 'datasets/rcs_confirm'

# 创建输出目录
os.makedirs(output_dir, exist_ok=True)

# 随机颜色映射（最多支持20个类别）
COLORS = [(int(255 * i / 20), int(255 * (20 - i) / 20), 128) for i in range(20)]

# 遍历标注文件
for label_file in os.listdir(labels_dir):
    if not label_file.endswith('.txt'):
        continue

    label_path = os.path.join(labels_dir, label_file)
    image_name = label_file.replace('.txt', '.jpg')  # 可根据实际扩展名修改
    image_path = os.path.join(images_dir, image_name)

    if not os.path.exists(image_path):
        print(f'Image {image_name} not found for label {label_file}')
        continue

    # 加载图像
    image = cv2.imread(image_path)
    h_img, w_img = image.shape[:2]

    # 读取标签
    with open(label_path, 'r') as f:
        lines = f.readlines()

    for line in lines:
        parts = line.strip().split()
        if len(parts) != 5:
            continue
        class_id, cx, cy, w, h = map(float, parts)
        class_id = int(class_id)

        # 转换为绝对坐标
        x1 = int((cx - w / 2) * w_img)
        y1 = int((cy - h / 2) * h_img)
        x2 = int((cx + w / 2) * w_img)
        y2 = int((cy + h / 2) * h_img)

        # 绘制矩形框和类名
        color = COLORS[class_id % len(COLORS)]
        cv2.rectangle(image, (x1, y1), (x2, y2), color, 2)
        cv2.putText(image, str(class_id), (x1, y1 - 5),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2)

    # 保存结果图像
    save_path = os.path.join(output_dir, image_name)
    cv2.imwrite(save_path, image)
    print(f'Saved {save_path}')
